package com.bw.test5;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.pipeline.Pipeline;
import com.alibaba.alink.pipeline.dataproc.vector.VectorAssembler;
import com.alibaba.alink.pipeline.feature.OneHotEncoder;
import com.alibaba.alink.pipeline.regression.LinearRegression;
import org.apache.flink.types.Row;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class PipelineTest {
    @Test
    public void testLinearRegression() throws Exception {
        BatchOperator.setParallelism(1);
        //1、构建数据集
        List <Row> df = Arrays.asList(
                Row.of(2, "a", 1),
                Row.of(3, "b", 1),
                Row.of(4, "c", 2),
                Row.of(2, "d", 1),
                Row.of(2, "b", 1),
                Row.of(4, "c", 2),
                Row.of(1, "b", 1)
        );
        BatchOperator <?> batchData = new MemSourceBatchOp(df, "f0 int, f1 string, label int");
        // 编码
        OneHotEncoder one_hot = new OneHotEncoder().setSelectedCols("f1").setOutputCols("f11");
        // 向量聚合
        VectorAssembler res = new VectorAssembler()
                .setSelectedCols("f0", "f11")
                .setOutputCol("vec");

        //2、创建算法组件
        LinearRegression lr = new LinearRegression()
                .setVectorCol("vec")
                .setLabelCol("label")
                .setPredictionCol("pred")
                .enableLazyPrintModelInfo()//延迟打印模型信息
                ;
        //3、训练
        //4、预测
        Pipeline pipeline = new Pipeline().add(one_hot).add(res).add(lr);

        BatchOperator<?> result = pipeline.fit(batchData).transform(batchData);

        result.print();

    }
}
